package com.inifitness.common.tools;

import java.util.Map;
import java.util.concurrent.Callable;

import lombok.AccessLevel;
import lombok.NoArgsConstructor;

/**
 * 多线程MDC工具
 *
 * @author sunjinfeng
 * @version 1.0
 * @since 2024/1/22 13:14
 */
@NoArgsConstructor(access = AccessLevel.PRIVATE)
public class ThreadMdcUtil {

  public static void setTraceIdIfAbsent() {
    if (MdcUtil.getTraceId() == null) {
      MdcUtil.setTraceId(MdcUtil.generateTraceId());
    }
  }

  public static <T> Callable<T> wrap(final Callable<T> callable,
      final Map<String, String> context) {
    return () -> {
      if (context == null) {
        MdcUtil.clear();
      } else {
        MdcUtil.setContextMap(context);
      }
      setTraceIdIfAbsent();
      try {
        return callable.call();
      } finally {
        MdcUtil.clear();
      }
    };
  }

  public static Runnable wrap(final Runnable runnable, final Map<String, String> context) {
    return () -> {
      if (context == null) {
        MdcUtil.clear();
      } else {
        MdcUtil.setContextMap(context);
      }
      //设置traceId
      setTraceIdIfAbsent();
      try {
        runnable.run();
      } finally {
        MdcUtil.clear();
      }
    };
  }
}
